from langgraph.graph import StateGraph, END
from typing import TypedDict, Optional, List
import asyncio
from ..strategy_agent.core.base_strategy import BacktestResult

class WorkflowState(TypedDict):
    prompt: str
    strategy_code: Optional[str]
    backtest_result: Optional[BacktestResult]
    optimization_history: List[dict]

class StrategyWorkflow:
    def __init__(self, llm_client, validator, backtest_engine):
        self.workflow = StateGraph(WorkflowState)
        self.llm = llm_client
        self.validator = validator
        self.backtest_engine = backtest_engine
        
        # 定义工作流节点
        self.workflow.add_node("brainstorm", self.brainstorm_ideas)
        self.workflow.add_node("generate_code", self.generate_strategy_code)
        self.workflow.add_node("execute_backtest", self.execute_backtest)
        self.workflow.add_node("optimize_params", self.optimize_parameters)
        
        # 设置流程路由
        self.workflow.set_entry_point("brainstorm")
        self.workflow.add_edge("brainstorm", "generate_code")
        self.workflow.add_conditional_edges(
            "generate_code",
            self.validate_code,
            {
                "valid": "execute_backtest",
                "invalid": "generate_code"
            }
        )
        self.workflow.add_conditional_edges(
            "execute_backtest",
            self.evaluate_backtest,
            {
                "acceptable": END,
                "needs_optimization": "optimize_params",
                "failed": "brainstorm"
            }
        )
        self.workflow.add_edge("optimize_params", "execute_backtest")
        
    async def brainstorm_ideas(self, state: WorkflowState) -> WorkflowState:
        """使用LLM生成策略思路"""
        prompt = f"生成量化交易策略，要求：{state['prompt']}"
        response = await self.llm.generate(prompt)
        return {**state, "strategy_idea": response}
    
    async def generate_strategy_code(self, state: WorkflowState) -> WorkflowState:
        """生成可执行策略代码"""
        code_prompt = f"""基于以下策略思路生成Python代码：
        {state['strategy_idea']}
        要求：
        - 继承BaseStrategy类
        - 实现required_factors和generate_signal方法
        - 包含合理的参数配置"""
        code = await self.llm.generate(code_prompt)
        return {**state, "strategy_code": code}
    
    async def validate_code(self, state: WorkflowState) -> str:
        """验证生成代码的合法性"""
        is_valid, errors = await self.validator.validate(state["strategy_code"])
        return "valid" if is_valid else "invalid"
    
    async def execute_backtest(self, state: WorkflowState) -> WorkflowState:
        """执行分布式回测"""
        strategy_class = self.validator.load_strategy(state["strategy_code"])
        result = await self.backtest_engine.run_backtest(
            strategy_class=strategy_class,
            params={},
            symbols=["000001.SH"],
            start="2020-01-01",
            end="2023-12-31"
        )
        return {**state, "backtest_result": result}
    
    async def evaluate_backtest(self, state: WorkflowState) -> str:
        """评估回测结果"""
        result = state["backtest_result"]
        if result.sharpe >= 1.5 and result.max_drawdown < 0.2:
            return "acceptable"
        elif result.sharpe >= 1.0:
            return "needs_optimization"
        else:
            return "failed"
    
    async def optimize_parameters(self, state: WorkflowState) -> WorkflowState:
        """贝叶斯参数优化"""
        optimization_config = {
            "parameters": {
                "window_size": {"min": 10, "max": 50, "type": "int"},
                "threshold": {"min": 0.01, "max": 0.1}
            },
            "max_iterations": 20
        }
        best_params = await self.backtest_engine.optimize(
            state["strategy_code"], 
            optimization_config
        )
        state["optimization_history"].append(best_params)
        return state

def create_workflow(llm_client, validator, backtest_engine) -> StateGraph:
    """创建工作流实例"""
    return StrategyWorkflow(llm_client, validator, backtest_engine).workflow